package com.atguigu.bigdata.sparkSql

import org.apache.spark.{SparkConf}
import org.apache.spark.sql.{Row, SparkSession}
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.Encoder
import org.apache.spark.sql.Encoders
import org.apache.spark.sql.SparkSession

//自定义聚合函数
object SparkSQL05_UDAF {
  def main(args: Array[String]): Unit = {
    //sparkConf
    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("SparkSQL01")
    //sparkSession
    //    val session:SparkSession = new SparkSession(sparkConf)
    val spark1 = SparkSession.builder().config(sparkConf).getOrCreate()
    import spark1.implicits._
    //创建聚合函数兑现那个
    val udaf = new MyAgeAvgFunction
//    将聚合函数转换为查询列
    val avgCol = udaf.toColumn.name("avgAge")
//
    var frame1 = spark1.read.json("in/name.json")
    val userDs = frame1.as[UserData]
    userDs.select(avgCol).show()
    spark1.stop()
  }

}

case class UserData(name: String, age: BigInt)

case class AvgBuffer(var sum: BigInt, var count:BigInt)

//  声明用户自定义函数（强类型）  Aggregator[-IN, BUF, OUT]
class MyAgeAvgFunction extends Aggregator[UserData, AvgBuffer, Double] {
  //  缓冲器初始化
  override def zero: AvgBuffer = {
    AvgBuffer(0, 0)
  }
//聚合数据
  override def reduce(b: AvgBuffer, a: UserData): AvgBuffer = {
    b.sum = b.sum + a.age
    b.count = b.count + 1
//  返回b
    b
  }
//缓冲区的合并操作
  override def merge(b1: AvgBuffer, b2: AvgBuffer): AvgBuffer = {
    b1.sum = b1.sum + b2.sum
    b1.count = b1.count + b2.count
    b1
  }
//完成计算
  override def finish(reduction: AvgBuffer): Double = {
    reduction.sum.toDouble / reduction.count.toDouble
  }

  override def bufferEncoder: Encoder[AvgBuffer] = Encoders.product

  override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}